# import torch.nn as nn
# import torch.nn.functional as F


# class RNNAgent(nn.Module):
#     def __init__(self, input_shape, args):
#         super(RNNAgent, self).__init__()
#         self.args = args

#         self.fc1 = nn.Linear(input_shape, args.rnn_hidden_dim)
#         self.rnn = nn.GRUCell(args.rnn_hidden_dim, args.rnn_hidden_dim)
#         self.fc2 = nn.Linear(args.rnn_hidden_dim, args.n_actions)

#     def init_hidden(self):
#         # make hidden states on same device as model
#         return self.fc1.weight.new(1, self.args.rnn_hidden_dim).zero_()

#     def forward(self, inputs, hidden_state=None):
#         b, a, e = inputs.size()
        
#         x = F.relu(self.fc1(inputs.view(-1, e)), inplace=True)
#         if hidden_state is not None:
#             hidden_state = hidden_state.reshape(-1, self.args.rnn_hidden_dim)
#         h = self.rnn(x, hidden_state)
#         q = self.fc2(h)
#         return q.view(b, a, -1), h.view(b, a, -1)

import torch as th
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import LayerNorm


class RNNAgent(nn.Module):
    def __init__(self, input_shape, args):
        super(RNNAgent, self).__init__()
        self.args = args

        # self.hidden_sizes = getattr(args, "hidden_sizes", [256, 256])
        self.rnn_hidden_dim = args.rnn_hidden_dim
        
        # self.layer_norm0 = LayerNorm(input_shape)
        self.fc1 = nn.Linear(input_shape, self.rnn_hidden_dim)
        # self.layer_norm1 = LayerNorm(self.hidden_sizes[0])
        self.rnn = nn.GRUCell(self.rnn_hidden_dim, self.rnn_hidden_dim)
        # self.layer_norm2 = LayerNorm(self.hidden_sizes[1])
        self.fc2 = nn.Linear(self.rnn_hidden_dim, args.n_actions)

        # initialization
        with th.no_grad():
            th.nn.init.orthogonal_(self.fc1.weight)
            th.nn.init.orthogonal_(self.fc2.weight, gain=self.args.gain)

    def init_hidden(self):
        # make hidden states on same device as model
        return self.fc1.weight.new(1, self.rnn_hidden_dim).zero_()

    def forward(self, x, hidden_state):
        # x = self.layer_norm0(x)
        x = F.relu(self.fc1(x), inplace=True)
        # x = self.layer_norm1(x)
        h_in = hidden_state.reshape(-1, self.rnn_hidden_dim)
        hidden_state = self.rnn(x, h_in)
        # x = self.layer_norm2(x)
        pi = self.fc2(hidden_state)

        return pi, hidden_state